package com.xjzd.std.stream;

import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

public class MultiThreadStream {

    static ExecutorService executorService = Executors.newFixedThreadPool(10);

    public static void main(String[] args) throws ExecutionException, InterruptedException {
        List<Student> oldStudentList = getStudentList();

        long l1 = System.currentTimeMillis();
        CompletableFuture<List<NewStudent>> newStudentList =
                pageProcessThenMerge(oldStudentList, Student::getId, MultiThreadStream::getNewStudentList, 2);
        System.out.println(newStudentList.get().size());
        System.out.println(System.currentTimeMillis() - l1);
        executorService.shutdown();
    }

    private static <T, TT, R> CompletableFuture<List<R>> pageProcessThenMerge(
            List<T> list, Function<T, TT> t2ttFun, Function<List<TT>, List<R>> tt2rFun, int pageSize
    ) {
        return CompletableFuture.supplyAsync(
                () -> IntStream.range(0, list.size()).boxed()
                        .collect(Collectors.groupingBy(i -> i / pageSize
                                , Collectors.mapping(list::get,
                                        Collectors.mapping(t2ttFun, Collectors.toList())
                                )))
                        .values().parallelStream()
                        .map(v -> CompletableFuture.supplyAsync(() -> tt2rFun.apply(v), executorService))
                        .collect(Collectors.collectingAndThen(
                                Collectors.toList(),
                                fList -> fList.stream().map(CompletableFuture::join).flatMap(Collection::stream).collect(Collectors.toList())
                        ))
                , executorService
        );
    }

    // 错误写法
    private static <T, TT, R> CompletableFuture<List<R>> pageProcessThenMergeError(
            List<T> list, Function<T, TT> t2ttFun, Function<List<TT>, List<R>> tt2rFun, int pageSize
    ) {
        return CompletableFuture.supplyAsync(
                () -> IntStream.range(0, list.size()).boxed()
                        .collect(Collectors.groupingBy(i -> i / pageSize
                                , Collectors.mapping(list::get,
                                        Collectors.mapping(t2ttFun, Collectors.toList())
                                )))
                        .values().parallelStream()
                        .map(v -> CompletableFuture.supplyAsync(() -> tt2rFun.apply(v), executorService))
                        .map(CompletableFuture::join)
                        .flatMap(List::stream)
                        .collect(Collectors.toList())
                , executorService
        );
    }

    private static List<Student> getStudentList() {
        return Arrays.asList(new Student(1L), new Student(2L), new Student(3L), new Student(4L), new Student(5L));
    }

    private static List<NewStudent> getNewStudentList(List<Long> idList) {
        try {
            Thread.sleep(idList.size() * 1000L);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        return idList.stream().flatMap(a -> Stream.of(new NewStudent(), new NewStudent())).collect(Collectors.toList());
    }

    static class Student {

        private Long id;

        public Student(Long id) {
            this.id = id;
        }

        public long getId() {
            return id;
        }

        public void setId(Long id) {
            this.id = id;
        }
    }

    static class NewStudent {

    }
}
